import os

import openai
from langchain.agents import AgentType, initialize_agent
from langchain.tools import Tool

from a0base.base_llm import pop_llm

openai.api_key = os.getenv("OPENAI_API_KEY")

# 1. 先创建一个「被调用」的 Agent
sub_agent = initialize_agent(
    tools=[],  # 它自己不需要工具
    llm=pop_llm,
    agent=AgentType.OPENAI_FUNCTIONS,
    description="擅长把用户问题改写为一句俏皮话"
)


# 2. 用 Tool 包装它，供主 Agent 调用
def ask_sub_agent(query: str) -> str:
    return sub_agent.run(query)


sub_tool = Tool(
    name="Joker",
    func=ask_sub_agent,
    description="一句话冷笑话生成器，输入任意主题即可"
)

# 3. 主 Agent
main_agent = initialize_agent(
    tools=[sub_tool],
    llm=pop_llm,
    agent=AgentType.OPENAI_FUNCTIONS
)

if __name__ == "__main__":
    print(main_agent.run("请用 1 句话介绍『量子计算』，要冷笑话风格"))
